#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import math
import numpy as np
import matplotlib.pyplot as plt
import csv
from pathlib import Path

# ---------- 调度函数 ----------
def _cos(x: float) -> float:
    """0→1 的余弦升温（带夹紧）"""
    x = min(1.0, max(0.0, x))
    return 0.5 * (1.0 - math.cos(math.pi * x))

def mask_schedule(step: int, total_steps: int):
    """
    返回 (p, L): 起点概率 p、span_len L（int）
    三阶段：0-0.3T 涨 p；0.3T-0.6T 涨 L；0.6T-1.0T 再涨 p
    超过 total_steps 后保持 p、L 不变
    """
    T = max(1, int(total_steps))
    s = max(0, int(step))
    # 超过总步数后钉死
    if s > T:
        s = T

    # 阶段边界
    a, b = int(0.3 * T), int(0.6 * T)
    a = max(1, a)
    b = max(a + 1, b)

    # 目标区间
    p0, p1_mid, p1 = 0.01, 0.05, 0.08
    L0, L1 = 3, 10

    if s <= a:
        alpha = _cos(s / a)
        p = p0 + (p1_mid - p0) * alpha
        L = L0
    elif s <= b:
        beta = _cos((s - a) / (b - a))
        p = p1_mid
        L = int(round(L0 + (L1 - L0) * beta))
    else:
        gamma = _cos((s - b) / (T - b))
        p = p1_mid + (p1 - p1_mid) * gamma
        L = L1

    # 安全夹紧
    p = float(min(0.99, max(0.0, p)))
    L = max(1, int(L))
    return p, L

def expected_coverage(p: float, L: int) -> float:
    """期望覆盖率的近似：1 - (1 - p)^L"""
    return 1.0 - (1.0 - p) ** L

# ---------- 参数 ----------
TOTAL_STEPS = 100_000      # 调度的最大步数
PLOT_END    = 120_000      # 绘制到 12 万步（最后 2 万应保持不变）
STRIDE      = 100          # 每 100 步取一个点

# ---------- 计算曲线 ----------
steps = np.arange(0, PLOT_END + 1, STRIDE, dtype=int)
ps, Ls, covs = [], [], []
for s in steps:
    p, L = mask_schedule(s, TOTAL_STEPS)
    ps.append(p)
    Ls.append(L)
    covs.append(expected_coverage(p, L))

# 打印关键点（100k 与 120k）
def at(step):
    p, L = mask_schedule(step, TOTAL_STEPS)
    cov = expected_coverage(p, L)
    return p, L, cov

print("At 100k :", at(100_000))
print("At 120k :", at(120_000))

# # ---------- 导出 CSV ----------
# out_csv = Path("mask_schedule_curve.csv")
# with out_csv.open("w", newline="") as f:
#     w = csv.writer(f)
#     w.writerow(["step", "p", "L", "expected_coverage"])
#     for s, p, L, c in zip(steps, ps, Ls, covs):
#         w.writerow([s, p, L, c])
# print(f"CSV saved to: {out_csv.resolve()}")

# ---------- 画图 ----------
plt.figure(figsize=(10,5))
plt.plot(steps, ps)
plt.title("Mask start probability p vs step")
plt.xlabel("step"); plt.ylabel("p"); plt.grid(True)
plt.tight_layout()
plt.savefig("mask_schedule_p.png", dpi=160)

plt.figure(figsize=(10,5))
plt.plot(steps, Ls)
plt.title("Span length L vs step")
plt.xlabel("step"); plt.ylabel("L"); plt.grid(True)
plt.tight_layout()
plt.savefig("mask_schedule_L.png", dpi=160)

plt.figure(figsize=(10,5))
plt.plot(steps, covs)
plt.title("Expected coverage vs step")
plt.xlabel("step"); plt.ylabel("expected coverage"); plt.grid(True)
plt.tight_layout()
plt.savefig("mask_schedule_cov.png", dpi=160)

print("Figures saved: mask_schedule_p.png, mask_schedule_L.png, mask_schedule_cov.png")
